'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: document_service.py
* @Time: 2025/11/12
* @All Rights Reserve By Brtc
'''
import logging
import random
import time
from dataclasses import dataclass
from datetime import datetime
from uuid import UUID

from click.core import batch
from injector import inject
from redis import Redis
from sqlalchemy.sql.expression import desc, asc
from sqlalchemy.sql.functions import func

from internal.entity.dataset_entity import ProcessType, SegmentStatus, DocumentStatus
from internal.entity.upload_file_entity import ALLOWED_FILE_EXTENSION
from internal.task.document_task import build_documents, update_document_enabled
from internal.model import Document, Dataset, UploadFile, ProcessRule, Segment
from pkg.paginator.paginator import Paginator
from internal.exception.exception import NotFoundException, ForbiddenException, FailException
from internal.schema.document_schema import GetDocumentsWithPageReq
from pkg.sqlalchemy import SQLAlchemy
from .base_service import BaseService
from ..entity.cache_entity import LOCK_DOCUMENT_UPDATE_ENABLED, LOCK_EXPIRE_TIME
from ..lib.helper import datetime_to_timestamp


@inject
@dataclass
class DocumentService(BaseService):
    db:SQLAlchemy
    redis_client:Redis

    def get_latest_position(self, dataset_id:UUID)->int:
        """根据传递的知识库id 获取最新文档位置"""
        document = self.db.session.query(Document).filter(
            Document.dataset_id == dataset_id,
        ).order_by(desc("position")).first()
        return document.position if document else 0



    def create_documents(self,
                         dataset_id:UUID,
                         upload_file_ids:list[UUID],
                         process_type:str == ProcessType.AUTOMATIC,
                         rule:dict = None) -> tuple[list[Document], str]:
        """根据传递的信息创建文档列表并调用异步任务"""
        # todo:等地授权认证模块完成
        account_id = "e036e651-9eae-4b15-8dc3-b9d692121acf"
        #1、检测知识库权限
        dataset = self.get(Dataset, dataset_id)
        if dataset is None or str(dataset.account_id) != account_id:
            raise ForbiddenException("当前用户无该知识库权限或知识库不存在！！")
        #2、提取并校验文件权限与文件扩展
        upload_files = self.db.session.query(UploadFile).filter(
            UploadFile.account_id == account_id,
            UploadFile.id.in_(upload_file_ids)
        ).all()
        upload_files = [
            upload_file for upload_file in upload_files
            if upload_file.extension.lower() in ALLOWED_FILE_EXTENSION
        ]

        if len(upload_files) == 0:
            logging.warning(f"上传的文件列表未接到合法的文件:{account_id}, dataset_id:{dataset_id}, upload_file_id:{upload_file_ids}")
            raise FailException("暂未解析到合法的文件, 请检查后重试")
        #3、创建批次与处理规则并记录到数据库中
        batch = time.strftime("%Y%m%d%H%M%S") + str(random.randint(100000, 999999))
        process_rule = self.create(
            ProcessRule,
            account_id=account_id,
            dataset_id=dataset_id,
            mode = process_type,
            rule=rule
        )
        #4、获取当前知识库最新的文档位置
        position = self.get_latest_position(dataset_id)
        #5、循环遍历所有合法的上传文件列表并记录
        documents = []
        for upload_file in upload_files:
            position +=1
            document = self.create(
                Document,
                account_id=account_id,
                dataset_id=dataset_id,
                upload_file_id=upload_file.id,
                process_rule_id=process_rule.id,
                batch=batch,
                name=upload_file.name,
                position=position,
            )
            documents.append(document)

        #6、调用异步任务
        build_documents.delay([document.id for document in documents])

        #7、返回文档列表与 处理批次
        return documents, batch


    def get_document(self, dataset_id: UUID, document_id: UUID) -> Document:
        """根据传递的知识库id+文档id获取文档记录信息"""
        # todo:等待授权认证模块完成进行切换调整
        account_id = "e036e651-9eae-4b15-8dc3-b9d692121acf"
        document = self.get(Document, document_id)
        if document is None:
            raise NotFoundException("该文档不存在，请核实后重试")
        if document.dataset_id != dataset_id or str(document.account_id) != account_id:
            raise ForbiddenException("当前用户获取该文档，请核实后重试")
        return document

    def update_document(self, dataset_id: UUID, document_id: UUID, **kwargs) -> Document:
        """根据传递的知识库id+文档id，更新文档信息"""
        # todo:等待授权认证模块完成进行切换调整
        account_id = "e036e651-9eae-4b15-8dc3-b9d692121acf"
        document = self.get(Document, document_id)
        if document is None:
            raise NotFoundException("该文档不存在，请核实后重试")
        if document.dataset_id != dataset_id or str(document.account_id) != account_id:
            raise ForbiddenException("当前用户无权限修改该文档，请核实后重试")
        return self.update(document, **kwargs)

    def get_documents_with_page(
            self, dataset_id: UUID, req: GetDocumentsWithPageReq,
    ) -> tuple[list[Document], Paginator]:
        """根据传递的知识库id+请求数据获取文档分页列表数据"""
        # todo:等待授权认证模块完成进行切换调整
        account_id = "e036e651-9eae-4b15-8dc3-b9d692121acf"
        # 1.获取知识库并校验权限
        dataset = self.get(Dataset, dataset_id)
        if dataset is None or str(dataset.account_id) != account_id:
            raise NotFoundException("该知识库不存在，或无权限")
        # 2.构建分页查询器
        paginator = Paginator(db=self.db, req=req)
        # 3.构建筛选器
        filters = [
            Document.account_id == account_id,
            Document.dataset_id == dataset_id,
        ]
        if req.search_word.data:
            filters.append(Document.name.ilike(f"%{req.search_word.data}%"))
        # 4.执行分页并获取数据
        documents = paginator.paginate(
            self.db.session.query(Document).filter(*filters).order_by(desc("created_at"))
        )
        return documents, paginator


    def get_document_status(self, dataset_id: UUID, batch:str)->list[dict]:
        """根据传递的知识库id + 批处理号 获取文档的状态"""
        # todo:等待授权认证模块完成进行切换调整
        account_id = "e036e651-9eae-4b15-8dc3-b9d692121acf"
        #1、检测知识库权限
        dataset = self.get(Dataset, dataset_id)
        if dataset is None or str(dataset.account_id) != account_id:
            raise ForbiddenException("当前用户无该只是库权限或知识库不存在！！")
        #2、查询当前知识库下该批次的文档列表
        documents = self.db.session.query(Document).filter(
            Document.dataset_id == dataset_id,
            Document.batch == batch,
        ).order_by(asc("position")).all()

        if documents is None or len(documents) == 0:
            raise NotFoundException("该批次未发现文档,请核实后重试！！")
        #3、循环遍历文档列表提取文档的状态信息
        document_status = []
        for document in documents:
            #4、查询每个文档的总片段数和已构建的完成的片段数量
            segment_count = self.db.session.query(func.count(Segment.id)).filter(
                Segment.document_id == document.id,
            ).scalar()
            completed_segment_count = self.db.session.query(func.count(Segment.id)).filter(
                Segment.document_id == document.id,
                Segment.status == SegmentStatus.COMPLETED,
            ).scalar()

            upload_file = document.upload_file
            document_status.append({
                "id": document.id,
                "name": document.name,
                "size":upload_file.size,
                "extension": upload_file.extension,
                "mime_type": upload_file.mime_type,
                "position": document.position,
                "segment_count": segment_count,
                "completed_segment_count": completed_segment_count,
                "error": document.error,
                "status": document.status,
                "processing_started_at":datetime_to_timestamp(document.processing_started_at),
                "parsing_completed_at":datetime_to_timestamp(document.parsing_completed_at),
                "splitting_completed_at":datetime_to_timestamp(document.splitting_completed_at),
                "indexing_completed_at":datetime_to_timestamp(document.indexing_completed_at),
                "completed_at":datetime_to_timestamp(document.completed_at),
                "stopped_at":datetime_to_timestamp(document.stopped_at),
                "created_at":datetime_to_timestamp(document.created_at),
            })
        return document_status


    def update_document_enabled(self, dataset_id: UUID, document_id: UUID, enabled: bool)->Document:
        """根据传递的知识库id + 文档id 更新文档的启用状态, 同时会 异步的更新 weaviate 向量数据库中的数据"""
        # todo:等待授权认证模块完成进行切换调整
        account_id = "e036e651-9eae-4b15-8dc3-b9d692121acf"
        #1、获取文档并校验
        document = self.get(Document, document_id)
        if document is None:
            raise NotFoundException("该文档不存在, 请核实后重试！")
        if document.dataset_id != dataset_id or str(document.account_id) != account_id:
            raise ForbiddenException("当前用户无权限修改 该文档！！")
        #2、判断文档是否处于可以修改状态 ， 只有构建完成文档才能被修改 enabled
        if document.status != DocumentStatus.COMPLETED:
            raise FailException(f"文档状态目前正不处于可修改状态！")
        #3、判断启用状态是否正确应该与数据的状态相反
        if document.enabled == enabled:
            raise ForbiddenException(f"文档状态修改错误 目前已经是{enabled}状态")
        #4、获取更新文档启用状态缓存锁并查看是都上锁
        cache_key = LOCK_DOCUMENT_UPDATE_ENABLED.format(document_id=document.id)
        cache_result = self.redis_client.get(cache_key)
        if cache_result is not None:
            raise FailException(f"当前文旦正在被修改，请稍后重试！！")
        #5、修改文档的启用状态，并设置缓存键， 缓存时间为 600s
        self.update(
            document,
            enabled=enabled,
            disabled_at = None if enabled else datetime.now(),
        )
        self.redis_client.setex(cache_key, LOCK_EXPIRE_TIME, 1)

        #6、其启用异步任务完成后续操作
        update_document_enabled.delay(document.id)
        return document

















